package com.scott.distrubutelock.javaapi;

import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooKeeper;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

/**
 * 分布式锁
 *
 * @author Scott 2017/11/2 18:08
 */
public class DistributeLock {

    /**
     * 分布式锁根节点名称
     */
    private static final String ROOT_LOCKS = "/LOCKS";

    //zookeeper
    private ZooKeeper zookeeper;

    /**
     * 锁ID
     */
    private String lockId;

    private final static byte[] data = {1, 2};

    /**
     * 会话超时时间
     */
    private int sessionTimeOut;

    public DistributeLock() throws IOException, InterruptedException {
        this.zookeeper = ZookeeperClient.getInstance();
        this.sessionTimeOut = ZookeeperClient.getSessionTimeOut();
    }

    private  CountDownLatch countDownLatch = new CountDownLatch(1);

    /**
     * 获取锁
     *
     * @return boolean
     */
    public boolean lock() {
        try {
            //在根节点下创建锁节点
            lockId = zookeeper.create(ROOT_LOCKS + "/", data, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
            System.out.println(Thread.currentThread().getName() + "成功创建节点->" + lockId + " 开始去竞争锁");

            //获取当前节点下的所有子节点
            List<String> childrenNodes = zookeeper.getChildren(ROOT_LOCKS, true);

            //构建锁的完成路径
            SortedSet<String> sortedSet = new TreeSet<>();
            for(String child:childrenNodes){
                sortedSet.add(ROOT_LOCKS+"/"+child);
            }
            //获取最小的节点
            String first = sortedSet.first();


            //如果获取的第一个节点和当前的lockId，说明最小节点就是锁，直接返回true
            if (lockId.equals(first)) {
                System.out.println(Thread.currentThread().getName() + "->成功获取到锁，lock节点为：" + lockId);
                return true;
            }
            //获取比当前锁小的所有节点列表
            SortedSet<String> lessThanLockId = sortedSet.headSet(lockId);
            if (!lessThanLockId.isEmpty()) {
                //获取最后一个，也就是比他大一点的节点对其进行监控，监控是否被删除，如果删除就说明下一个就是他自己了，他就能到锁了
                String prevLockId = lessThanLockId.last();
                zookeeper.exists(prevLockId, new LockWatcher(countDownLatch));

                //需要添加一个countDownLatch的超时时间
                countDownLatch.await(sessionTimeOut, TimeUnit.MILLISECONDS);

                //如果上面代码执行完成了，则说明存在两种情况，要么监听到他上一个节点被删除，现在最小节点就是lockId,获取到锁，要么就是会话超时了
                System.out.println(Thread.currentThread().getName() + " 成功获取锁：" + lockId);
            }
            return true;


        } catch (KeeperException | InterruptedException e) {
            e.printStackTrace();
        }
        return false;
    }

    /**
     * 释放锁
     *
     * @return Boolean
     */
    public boolean unLock() {

        System.out.println(Thread.currentThread().getName() + "->开始释放锁：" + lockId);

        try {
            zookeeper.delete(lockId, -1);
            return true;
        } catch (InterruptedException | KeeperException e) {
            e.printStackTrace();
        }
        return false;
    }

    /**
     * 测试方法
     *
     * @param args
     */
    public static void main(String[] args) {
        final CountDownLatch latch = new CountDownLatch(10);
        Random random = new Random();
        for (int i = 0; i < 10; i++) {
            new Thread(() -> {
                DistributeLock lock = null;

                try {
                    lock = new DistributeLock();
                    latch.countDown();
                    latch.await();
                    lock.lock();
                    Thread.sleep(random.nextInt(500));

                } catch (IOException | InterruptedException e) {
                    e.printStackTrace();
                } finally {
                    if(lock!=null){
                        lock.unLock();
                    }
                }
            }).start();
        }
    }

}
